package com.bruce.ai.alibaba.workflow.config;

import java.util.HashMap;
import java.util.Map;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.bruce.ai.alibaba.workflow.node.ConditionEvaluatorNode;
import com.bruce.ai.alibaba.workflow.node.TextProcessorNode;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;


/**
 * 工作流配置类
 * 功能：配置工作流中的节点 Bean
 */
@Configuration
public class WorkflowConfiguration {

    @Bean
    public StateGraph customWorkflowGraph(ChatClient.Builder chatClientBuilder) throws Exception{
        //定义策略路线
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("query", new ReplaceStrategy());
            strategies.put("processed_text", new ReplaceStrategy());
            strategies.put("queryVariants", new ReplaceStrategy());
            strategies.put("final_result", new ReplaceStrategy());
            return strategies;
        };
        // 构建 StateGraph
        StateGraph graph = new StateGraph(keyStrategyFactory);

        // 添加自定义 Node
        graph.addNode("processor", node_async(new TextProcessorNode()));
        graph.addNode("condition", node_async(new ConditionEvaluatorNode()));

        // 添加边
        graph.addEdge(StateGraph.START, "processor");
        graph.addEdge("processor", "condition");

        // 条件边：根据 condition node 的结果路由
		graph.addConditionalEdges(
				"condition",
				edge_async(state -> state.value("_condition_result", "short").toString()),
				Map.of(
						"long", "processor",  // 长文本重新处理
						"short", StateGraph.END  // 短文本结束
				)
		);
        return graph;
    }

}
